try:
    from PyQt5.QtGui import QImage
except ImportError:
    from PyQt4.QtGui import QImage

try:
    from PyQt5.QtGui import *
    from PyQt5.QtCore import *
    from PyQt5.QtWidgets import *
except ImportError:
    # needed for py3+qt4
    # Ref:
    # http://pyqt.sourceforge.net/Docs/PyQt4/incompatible_apis.html
    # http://stackoverflow.com/questions/21217399/pyqt4-qtcore-qvariant-object-instead-of-a-string
    if sys.version_info.major >= 3:
        import sip
        sip.setapi('QVariant', 2)
    from PyQt4.QtGui import *
    from PyQt4.QtCore import *

import codecs
import json
import os
import sys
from pathlib import Path
from libs.constants import DEFAULT_ENCODING
from libs.pascal_voc_io import PascalVocWriter
from libs.yolo_io import YoloReader
import cv2


class annot_convert:
    def __init__(self):
        self.img_suffix = ['png', 'jpg']
        predef_classes_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data", "predefined_classes.txt")
        self.label_hist = list()
        if os.path.exists(predef_classes_file) is True:
            with codecs.open(predef_classes_file, 'r', 'utf8') as f:
                for line in f:
                    line = line.strip()
                    self.label_hist.append(line)

        return

    def read(self, filename, default=None):
        try:
            reader = QImageReader(filename)
            reader.setAutoTransform(True)
            return reader.read()
        except:
            return default

    def create_write(self, name, img_root):
        img_file_name = name[:name.rfind('.')]
        if os.path.exists(os.path.join(img_root, '{}.{}'.format(img_file_name, 'jpg'))):
            img_file_name = '{}.{}'.format(img_file_name, 'jpg')
        elif os.path.exists(os.path.join(img_root, '{}.{}'.format(img_file_name, 'png'))):
            img_file_name = '{}.{}'.format(img_file_name, 'png')
        else:
            print('{} not found in [png, jpg]'.format(img_file_name))
            return None

        image_path = os.path.join(img_root, img_file_name)
        img_folder_path = os.path.dirname(image_path)
        img_folder_name = os.path.split(img_folder_path)[-1]
        img_file_name = os.path.basename(image_path)

        image = self.read(image_path)
        image_shape = [image.height(), image.width(), 1 if image.isGrayscale() else 3]

        writer = PascalVocWriter(img_folder_name, img_file_name,
                                 image_shape, local_img_path=image_path)
        writer.verified = False
        return writer

    def txt2voc(self, txt_name, txt_path, writer):
        # out_file.write("%d %.6f %.6f %.6f %.6f\n" % (class_index, x_center, y_center, w, h))

        with open(os.path.join(txt_path, txt_name), mode='r') as r:
            height, width, c = writer.img_size

            while True:
                line = r.readline()
                if not line:
                    break
                line = line.strip('\n').strip(' ')
                arr = line.split(' ')
                class_index = int(arr[0])
                x_center = float(arr[1])
                y_center = float(arr[2])
                w = float(arr[3])
                h = float(arr[4])

                x_min = max(float(x_center) - float(w) / 2, 0)
                x_max = min(float(x_center) + float(w) / 2, 1)
                y_min = max(float(y_center) - float(h) / 2, 0)
                y_max = min(float(y_center) + float(h) / 2, 1)

                x_min = round(width * x_min)
                x_max = round(width * x_max)
                y_min = round(height * y_min)
                y_max = round(height * y_max)

                # x_min = x_center - w / 2
                # y_min = y_center - h / 2
                # x_max = x_center + w / 2
                # y_max = y_center + h / 2

                label = self.label_hist[class_index]
                # Add Chris
                difficult = 0
                writer.add_bnd_box(x_min, y_min, x_max, y_max, label, difficult)

            image = self.read(writer.local_img_path)
            t_yolo_parse_reader = YoloReader(os.path.join(txt_path, txt_name), image)
            shapes = t_yolo_parse_reader.get_shapes()

        return

    def convert2voc(self, convert_path, voc_path, img_root):
        os.makedirs(voc_path, exist_ok=True)

        for name in os.listdir(convert_path):
            if name.endswith('.json'):
                writer = self.create_write(name, img_root)

                line = Path(os.path.join(convert_path, name)).read_text(DEFAULT_ENCODING)
                annotation_list = json.loads(line)
                for annotations in annotation_list:
                    for annotation in annotations['annotations']:
                        label = annotation['label']
                        corrdinates = annotation['coordinates']
                        x = corrdinates['x']
                        y = corrdinates['y']
                        w = corrdinates['width']
                        h = corrdinates['height']

                        x_min = x - w / 2
                        y_min = y - h / 2
                        x_max = x + w / 2
                        y_max = y + h / 2
                        # Add Chris
                        difficult = 0
                        writer.add_bnd_box(x_min, y_min, x_max, y_max, label, difficult)

                xml_name = name[:name.rfind('.')] + '.xml'
                writer.save(os.path.join(voc_path, xml_name))

            elif name.endswith('txt'):
                writer = self.create_write(name, img_root)
                if writer is None:
                    continue

                self.txt2voc(name, convert_path, writer)
                xml_name = name[:name.rfind('.')] + '.xml'
                writer.save(os.path.join(voc_path, xml_name))
            else:
                print('the {} dont need convert'.format(name))
        return


if __name__ == '__main__':
    convert = annot_convert()
    convert.convert2voc(sys.argv[1], sys.argv[2], sys.argv[3])
